// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use anyhow::{bail, Context};
use std::{env, time::Duration};
use tracing::{debug, error, trace};

use xgpu_common::ipc::{
    error::IpcError,
    framer::LengthPrefixFramer,
    message::{Request, Response},
    peer::Client,
    transport::shmem::{ShmemTransport, ShmemTransportError},
};

mod api;
mod api_handler;
mod cache;
mod function_loader;

use std::fs::File;
use std::path::Path;
use tracing_subscriber::{fmt, prelude::*, EnvFilter, Registry};

fn init_tracing(file_path: Option<&str>) {
    let filter = EnvFilter::new("ERROR").add_directive("xgpu_common::ipc=error".parse().unwrap()); // ipc mod

    let mut layers = Vec::new();

    if let Some(path) = file_path {
        let file = File::create(Path::new(path)).expect("create log file failed");
        let file_layer = fmt::layer()
            .with_writer(file)
            .with_ansi(false)
            .with_filter(filter);
        layers.push(file_layer.boxed());
    } else {
        let stdout_layer = tracing_subscriber::fmt::layer()
            .with_ansi(true)
            .with_filter(filter.clone());
        layers.push(stdout_layer.boxed());
    }

    Registry::default().with(layers).init();
}

fn main() {
    init_tracing(None); // -> stdout

    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        eprintln!("Usage: {} <shmem addr>", args[0]);
        std::process::exit(1);
    }

    let mut threads = vec![];

    let pid = &args[1];
    for i in 0..32 {
        let addr = format!("{}_{}", pid, i);
        let handle = std::thread::spawn(|| self::server_main(addr));

        threads.push(handle);
    }

    for (tid, handle) in threads.into_iter().enumerate() {
        if let Err(e) = handle.join().expect("Failed to join server thread") {
            error!("Thread {}: {:?}", tid, e);
        }
    }
}

fn server_main(addr: String) -> anyhow::Result<()> {
    const BUFFER_SIZE: usize = 64 * 1024 * 1024;
    const CONNECT_TIMEOUT: Duration = Duration::from_secs(3);

    let transport = ShmemTransport::new(BUFFER_SIZE);
    let framer = LengthPrefixFramer::new(BUFFER_SIZE);

    let client = match Client::connect(&transport, framer, &addr, CONNECT_TIMEOUT) {
        Ok(c) => c,
        Err(IpcError::TransportError(ShmemTransportError::ConnectTimeout)) => {
            trace!("'{}': Connection timeout", addr);
            return Ok(());
        }
        Err(e) => {
            bail!("{}", e);
        }
    };
    debug!("{:#?}", client);

    loop {
        let mut request = match client.receive_message::<Request>() {
            Ok(Some(r)) => r,
            Ok(None) => continue,
            Err(IpcError::TransportError(ShmemTransportError::ConnectionClosed)) => {
                break;
            }
            Err(e) => bail!("{}", e),
        };
        debug!(
            "[Server] Received request: request_id={}, method_id={}, argc={}",
            request.request_id(),
            request.method_id(),
            request.argc()
        );

        let ret = api_handler::call_handler(request.method_id(), request.args_mut()).with_context(
            || {
                format!(
                    "Failed to handle request, request_id={}, method_id={}",
                    request.request_id(),
                    request.method_id()
                )
            },
        )?;

        let response = Response::with_request(&request, ret);
        client
            .send_message(&response)
            .context("Failed to send response")?;
    }

    Ok(())
}
